package com.example.aidemo.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.ai.zhipuai.ZhiPuAiImageModel;
import org.springframework.ai.zhipuai.ZhiPuAiImageOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.Map;

/**
 * 智谱 AI
 */
@RestController
@RequestMapping("/zp")
public class ZhipuAiController {

    @Autowired
    private ZhiPuAiChatModel chatModel;

    @Autowired
    private ZhiPuAiImageModel zhiPuAiImageModel;

    // 基于内存的聊天记忆
    private final ChatMemory chatMemory = new InMemoryChatMemory();

    /**
     * zhipuAi 聊天
     *
     * @param message
     * @return
     */
    @GetMapping("/ai/generate")
    public Map generate(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        return Map.of("generation", chatModel.call(message));
    }

    /**
     * 模型调用 stream 流式输出
     *
     * @param message
     * @return
     */
    @GetMapping(value = "/generateStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> generateStream(@RequestParam(value = "message", defaultValue = "糖醋排骨怎么做") String message) {
        //var prompt = new Prompt(new UserMessage(message));
        return chatModel.stream(message);
    }

    /**
     * 带有聊天记录功能的 ChatClient使用(流式调用)
     *
     * @param question
     * @param sessionId
     * @return
     */
    @GetMapping(value = "/chat", params = "question", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> getZhiPuAIMessage(String question, String sessionId) {
        // 使用InMemoryChatMemory进行内存存储记忆 sessionId根据id找对应的记录,只需要最近3000条
        MessageChatMemoryAdvisor messageChatMemoryAdvisor = new MessageChatMemoryAdvisor(chatMemory, sessionId, 3000);
        return ChatClient.create(chatModel).prompt().user(question)
                .advisors(messageChatMemoryAdvisor) // 查找记录一起发送给大模型
                .stream().content().map(
                        chatResponse -> ServerSentEvent.builder(chatResponse).event("message").build());
    }

    /**
     * 智普文生图
     *
     * @param message
     * @return
     */
    @GetMapping("/imageGeneration")
    public String imageGeneration(@RequestParam(value = "message", defaultValue = "画一幅向日葵的图片") String message) {
        ImageResponse response = zhiPuAiImageModel.call(
                new ImagePrompt(
                        message,
                        ZhiPuAiImageOptions.builder()
                                .withModel("cogview-3")
                                .build())
        );
        String url = response.getResult().getOutput().getUrl();
        // 设置整个页面的背景色为黑色，图片居中显示，高度适配屏幕
        return "<html><head><style>body { background-color: black;}</style><body><img src='" + url + "' style='display: block; margin: auto; height: 100vh;' /></body></html>";
    }

}
